import sys 
sys.path.append('code')
from feature_eng.alpha_huangeven.feature import Alphas_even_p1, Alphas_even_p0, Alphas_even_p2
import os  
from joblib import Parallel, delayed
import itertools
import pandas as pd
import time
import numpy as np

def calc_alpha_huangeven1(num='volume_relative_ratio5',param=5, columns_root=None, save_root=None, price_root=None, mp=60):
    t1 = time.time()
    if price_root is None:
        price_root = 'data/stock_data/daily/price'
    if columns_root is None:
        columns_root = r'data/stock_data/daily/consentrate_price/close.pkl.gzip'
    if save_root is None:
        save_root = r'data/cmodty/my_feature/feature_data_huangeven'
    if not os.path.exists(save_root):
        os.makedirs(save_root)
    data = pd.read_pickle(columns_root)
    column = data.columns
    lst = []
    for code in column:
        file = os.path.join(price_root,f'{code}.pkl.gzip')
        df_i = pd.read_pickle(file).astype(float)
        alpha_i = pd.DataFrame(index = df_i.index[mp:]) 
        alp1 = Alphas_even_p1(df_i, mp, param)
        ary = eval(f'alp1.{num}')()
        # print(ary.shape, df_i.shape, ary.shape[0]==df_i.shape[0]-mp )
        alpha_i[code] = ary
        lst.append(alpha_i)
    df = pd.concat(lst, axis=1)
    df.columns = column
    df = df[df.index>'2005-01-01']
    save_file = os.path.join(save_root, f'{num}_{param}.pkl.gzip')
    df.to_pickle(save_file)
    print(f"{num} calculating cost time:{(time.time()-t1):.2f}s")
    
def calc_alpha_huangeven2(num='volume_relative_ratio5',param1=5, param2=2, columns_root=None, save_root=None, price_root=None, mp=60):
    t1 = time.time()
    if price_root is None:
        price_root = 'data/stock_data/daily/price'
    if columns_root is None:
        columns_root = r'data/stock_data/daily/consentrate_price/close.pkl.gzip'
    if save_root is None:
        save_root = r'data/cmodty/my_feature/feature_data_huangeven'
    if not os.path.exists(save_root):
        os.makedirs(save_root)
    data = pd.read_pickle(columns_root)
    column = data.columns
    lst = []
    for code in column:
        file = os.path.join(price_root,f'{code}.pkl.gzip')
        df_i = pd.read_pickle(file).astype(np.float32)
        alpha_i = pd.DataFrame(index = df_i.index[mp:]) 
        alp2 = Alphas_even_p2(df_i, mp, param1, param2)
        ary =  eval(f'alp2.{num}')()
        # print(ary.shape, df_i.shape, ary.shape[0]==df_i.shape[0]-mp )
        alpha_i[code] = ary
        lst.append(alpha_i)
    df = pd.concat(lst, axis=1)
    df.columns = column
    save_file = os.path.join(save_root, f'{num}_{param1}_{param2}.pkl.gzip')
    df.to_pickle(save_file)
    print(f"{num} calculating cost time:{(time.time()-t1):.2f}s")
    
def calc_alpha_huangeven0(num='volume_relative_ratio5', columns_root=None, save_root=None, price_root=None, mp=60):
    t1 = time.time()
    if not os.path.exists(save_root):
        os.makedirs(save_root)
    data = pd.read_pickle(columns_root).astype(np.float32)
    column = data.columns
    lst = []
    for code in column:
        file = os.path.join(price_root,f'{code}.pkl.gzip')
        df_i = pd.read_pickle(file)
        alpha_i = pd.DataFrame(index = df_i.index[mp:]) 
        alp0 = Alphas_even_p0(df_i, mp)
        ary =  eval(f'alp0.{num}')()
        # print(ary.shape, df_i.shape, ary.shape[0]==df_i.shape[0]-mp )
        alpha_i[code] = ary
        lst.append(alpha_i)
    df = pd.concat(lst, axis=1)
    df.columns = column
    save_file = os.path.join(save_root, f'{num}_.pkl.gzip')
    df.to_pickle(save_file)
    print(f"{num} calculating cost time:{(time.time()-t1):.2f}s")

# def calc_etf_main(num):
#     columns_root=r'data\etf_data\consentrate_daily_price\close.pkl.gzip'
#     save_root=r'data/cmodty/etf_feature/feature_data_huangeven'
#     price_root=r'data\etf_data\daily_price'
#     calc_alpha_huangeven(num, columns_root, save_root, price_root)

def calc_main_even():
    p1_lst = ['volume_relative_ratio', 'close_dis', 'bias_calc', 'atr_calc', 'rsi_calc', 'cci_calc', 'roc_calc', 'trix_calc', 'mfi_calc', 'boll_up', 
                'boll_down', 'ar_calc', 'br_calc', 'vma_calc',  'open_logrt', 'rt_calc', 'close_logrt', 'log_volum', 'price_efficiency', 'amplitude']
    p2_lst = ['price_center', 'ma_dis']
    p0_lst = ['diffreturn','volume_change','macd','ptv_calc','obv0','obv1']
    price_root = 'data/stock_data/daily/price'
    columns_root = r'data/stock_data/daily/consentrate_price/close.pkl.gzip'
    save_root = r'data/cmodty/my_feature/feature_data_huangeven'
    param1_lst = [5,10,30,60] #5,10,30,
    Parallel(n_jobs=16)(delayed(calc_alpha_huangeven1)(i,j,columns_root, save_root, price_root, mp=240) for i,j in itertools.product(p1_lst, param1_lst))
    Parallel(n_jobs=16)(delayed(calc_alpha_huangeven0)(i,columns_root, save_root, price_root, mp=240) for i in p0_lst)
    Parallel(n_jobs=16)(delayed(calc_alpha_huangeven2)(i,j,k,columns_root, save_root, price_root, mp=240) for i,j,k in itertools.product(p2_lst, param1_lst, param1_lst))
    
    
# calc_main_even()